This PR replaces all occurrences of torch.cuda.amp.GradScaler(args...) and torch.cuda.amp.autocast(args...) with torch.GradScaler("cuda", args...) and torch.autocast("cuda", args...) respectively.

This fixes the deprecation warnings during training and inference like:

FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast("cuda", args...)` instead.
  with torch.cuda.amp.autocast(enabled=self.scaler is not None):

see PyTorch Docs

P.S. I'm not sure about the difference between torch.autocast("cuda", args...) and torch.amp.autocast("cuda", args...) (as seen in the deprecation warning). But I would to prefer to stick to the Documentation. E.g. only using torch.autocast without the amp